// new version of getliveip, fast detect live IP
// xiangyang, 2012.09.29
// Should be compiled as root:
//  0. su
//  1. go build getliveip.go
//  2. chmod u+s getliveip

package main

import (
  "net"
  "os"
  //"math/rand"
  "math/big"
  "time"
  "fmt"
  "strings"
  "strconv"
  //"runtime"
)

const payload = "This is a probe from argus.csnet1.cs.tsinghua.edu.cn"

func main() {
  if len(os.Args) != 2 {
    fmt.Fprintln(os.Stderr, "Usage:", os.Args[0], "prefix")
    os.Exit(1)
  }
  //runtime.GOMAXPROCS(runtime.NumCPU())
  //rand.Seed(time.Now().Unix())
  POW_GO_ROUTINE := uint32(10) // 2^POW_GO_ROUTINE is the number of go routine

  // parse IP prefix
  addr := make([]string, 0)
  maskLen, _ := strconv.Atoi(os.Args[1][strings.Index(os.Args[1], "/")+1 : ])
  netLen := uint32(0)
  isIPv6 := false
  if strings.Index(os.Args[1], ":") != -1 {
    netLen = uint32(128 - maskLen)
    isIPv6 = true
    POW_GO_ROUTINE -= 1 // IPv6 packet is bigger
  } else {
    netLen = uint32(32 - maskLen)
  }
  _, network, _ := net.ParseCIDR(os.Args[1])
  ip := network.IP
  oneBig := big.NewInt(1)
  ipNum := big.NewInt(0)
  ipNum.SetBytes(ip)
  ipNumBytes := make([]byte, 0)
  // gateway
  ipNum.Add(ipNum, oneBig)
  ipNumBytes = ipNum.Bytes()
  for j := 0; j < len(ipNumBytes); j++ {
    ip[j] = ipNumBytes[j]
  }
  gateway := ip.String()

  // IP address candidates in prefix
  if netLen <= POW_GO_ROUTINE { // if prefix is small, choose every IP
    for i := 1; i < (1 << netLen) - 1; i++ {
      ipNumBytes = ipNum.Bytes()
      for j := 0; j < len(ipNumBytes); j++ {
        ip[j] = ipNumBytes[j]
      }
      addr = append(addr, ip.String())
      ipNum.Add(ipNum, oneBig)
    }
  } else { // else, choose gateway (the first and last IP) of every subnet
    threeBig := big.NewInt(3)
    subSizeBig := big.NewInt(0)
    subSizeBig.Exp(
      big.NewInt(2), big.NewInt(int64(netLen - POW_GO_ROUTINE + 1)), nil)
    subSizeBig.Sub(subSizeBig, threeBig)
    for i := 0; i < (1 << (POW_GO_ROUTINE - 1)); i++ {
      // first IP in subnet
      ipNumBytes = ipNum.Bytes()
      for j := 0; j < len(ipNumBytes); j++ {
        ip[j] = ipNumBytes[j]
      }
      addr = append(addr, ip.String())
      // last IP in subnet
      ipNum.Add(ipNum, subSizeBig)
      ipNumBytes = ipNum.Bytes()
      for j := 0; j < len(ipNumBytes); j++ {
        ip[j] = ipNumBytes[j]
      }
      addr = append(addr, ip.String())
      // next subnet
      ipNum.Add(ipNum, threeBig)
    }
  }
  //fmt.Println(len(addr))
  //for i := 0; i < len(addr) && i < 10; i++ {
  //  fmt.Println(addr[i])
  //}//DEBUG
  //for i := len(addr) - 10; i < len(addr); i++ {
  //  fmt.Println(addr[i])
  //}//DEBUG

  // timeout channel
  timeout := make (chan bool, 1)
  go func () {
    time.Sleep(2e9)
    timeout <- true
  }()

  //// partial checksum
  //var msg [80]byte
  //if isIPv6 {
  //  msg[0] = 128 // type: echo 118
  //} else {
  //  msg[0] = 8 // type: echo 8
  //}
  //msg[1] = 0   // code: 0
  //msg[2] = 0   // checksum
  //msg[3] = 0   // checksum
  //for i := 0; i < len(payload); i++ {
  //  msg[8 + i] = byte(payload[i])
  //}
  //mlen := 8 + len(payload)
  //partSum := uint32(0)
  //if isIPv6 {
  //  protocol := "ip6:58" //ICMPv6
  //  conn, err := net.DialTimeout(protocol, gateway, 2 * time.Second)
  //  if err != nil {
  //    fmt.Println(gateway, "0")
  //    return
  //  }
  //  defer conn.Close()
  //  partSum = partialCheckSum(msg[0:mlen],
  //      net.ParseIP(conn.LocalAddr().String()))
  //} else {
  //  partSum = partialCheckSum(msg[0:mlen], nil)
  //}

  // live IP channel
  liveIP := make (chan string, 1)
  for i := 0; i < len(addr); i++ {
    go sendIcmp(uint32(i), uint32(0), addr[i], isIPv6, liveIP)
  }

  // get live IP
  select {
    case ip := <-liveIP:
      fmt.Println(ip, "1")
    case <-timeout:
      fmt.Println(gateway, "0")
  }

  os.Exit(0)
}

func sendIcmp(iden, partSum uint32, addr string, isIPv6 bool, liveIP chan string) {
  //identifier := uint32(rand.Int()) & 0xffff
  //sequence := uint32(rand.Int()) & 0xffff
  identifier := iden & 0xffff
  sequence := iden & 0xffff

  // packet
  var msg [128]byte
  if isIPv6 {
    msg[0] = 128 // type: echo 128
  } else {
    msg[0] = 8 // type: echo 8
  }
  msg[1] = 0   // code: 0
  msg[2] = 0   // checksum
  msg[3] = 0   // checksum
  msg[4] = byte((identifier >> 8) & 0xff) // identifier[0]
  msg[5] = byte(identifier & 0xff)        // identifier[1]
  msg[6] = byte((sequence >> 8) & 0xff)   // sequence[0]
  msg[7] = byte(sequence & 0xff)          // sequence[1]
  for i := 0; i < len(payload); i++ {
    msg[8 + i] = byte(payload[i])
  }
  mlen := 8 + len(payload)

  // connection
  protocol := "ip4:icmp"
  if isIPv6 {
    protocol = "ip6:58" //ICMPv6
  }
  conn, err := net.DialTimeout(protocol, addr, 2 * time.Second)
  if err != nil {
    //fmt.Fprintf(os.Stderr, "Fatal error: %s", err.Error())//DEBUG
    return
  }
  defer conn.Close()

  // checksum
  check := uint16(0)
  if isIPv6 {
    check = checkSum(msg[:mlen],
      net.ParseIP(conn.LocalAddr().String()),
      net.ParseIP(addr))
    //check = completeCheckSum(partSum, msg[:mlen], net.ParseIP(addr))
  } else {
    check = checkSum(msg[:mlen], nil, nil)
    //check = completeCheckSum(partSum, msg[:mlen], nil)
  }
  msg[2] = byte(check >> 8)
  msg[3] = byte(check & 0xff)

  _, err = conn.Write(msg[:mlen])
  if err != nil {
    return
  }
  //fmt.Println("Msg     ", msg[:mlen]) // DEBUG
  _, err = conn.Read(msg[:])
  if err != nil {
    return
  }
  //fmt.Println("Response", msg[:mlen]) // DEBUG

  if (isIPv6 && msg[0] != 129) || (!isIPv6 && msg[0] != 0) {
    return
  }

  if (uint32(msg[4]) << 8) | uint32(msg[5]) == identifier &&
     (uint32(msg[6]) << 8) | uint32(msg[7]) == sequence {
    liveIP <- addr
  }
}

func partialCheckSum(msg, localAddr []byte) uint32 {
  sum := uint32(0)

  if localAddr != nil { // ICMPv6
    //fmt.Println(len(localAddr), localAddr)
    for i := 0; i < 16; i += 2 {
      sum += (uint32(localAddr[i]) << 8) | uint32(localAddr[i+1])
    }
    sum += uint32(len(msg))
    sum += 58
  }

  for n := 0; n < len(msg) - 1; n += 2 {
    if n >= 4 || n <= 7 {
      continue
    } // identifier & sequence
    sum += (uint32(msg[n]) << 8) | uint32(msg[n+1])
  }
  if len(msg) & 1 != 0 {
    sum += uint32(msg[len(msg) - 1])
  }

  return sum
}

func completeCheckSum(partSum uint32, msg, remoteAddr []byte) uint16 {
  if remoteAddr != nil { // ICMPv6
    //fmt.Println(len(remoteAddr), remoteAddr)
    for i := 0; i < 16; i += 2 {
      partSum += (uint32(remoteAddr[i]) << 8) | uint32(remoteAddr[i+1])
    }
  }

  for n := 4; n <= 7; n += 2 { // identifier & sequence
    partSum += (uint32(msg[n]) << 8) | uint32(msg[n+1])
  }

  partSum = (partSum >> 16) + (partSum & 0xffff)
  partSum += (partSum >> 16)
  return uint16(^partSum)
}

func checkSum(msg, localAddr, remoteAddr []byte) uint16 {
  sum := uint32(0)

  if localAddr != nil && remoteAddr != nil { // ICMPv6
    //fmt.Println(len(localAddr), localAddr)
    //fmt.Println(len(remoteAddr), remoteAddr)
    for i := 0; i < 16; i += 2 {
      sum += (uint32(localAddr[i]) << 8) | uint32(localAddr[i+1])
    }
    for i := 0; i < 16; i += 2 {
      sum += (uint32(remoteAddr[i]) << 8) | uint32(remoteAddr[i+1])
    }
    sum += uint32(len(msg))
    sum += 58
  }

  for n := 0; n < len(msg) - 1; n += 2 {
    sum += (uint32(msg[n]) << 8) | uint32(msg[n+1])
  }
  if len(msg) & 1 != 0 {
    sum += uint32(msg[len(msg) - 1])
  }

  sum = (sum >> 16) + (sum & 0xffff)
  sum += (sum >> 16)
  return uint16(^sum)
}

